# encoding: utf-8
# encoding: utf-8

from typing import List
from pydantic import BaseModel
from fastapi import FastAPI
import uvicorn

from named_entity import NamedEntityExtractor

pretrained_model = "E:/nlp_about/pretrained_model/nlp_raner_named-entity-recognition_chinese-base-ecom-50cls"

extractor = NamedEntityExtractor(pretrained_model)

# sku = "得力 0010 订书钉"
#
# print(extractor.extract(sku))
app = FastAPI()


class Item(BaseModel):
    text: str = None
    texts: List[str] = None


@app.post("/ner-api", tags=[])
def post_ner(item: Item):
    if not item.text and not item.texts:
        return {}
    if item.text:
        return extractor.extract(item.text)
    if item.texts:
        return extractor.extract_batch(item.texts)


@app.get("/ner/{text}", tags=[])
def get_ner(text: str):
    if not text:
        return {}
    return extractor.extract(text)


if __name__ == '__main__':
    uvicorn.run(app, host="0.0.0.0", port=10001)
